#!/usr/bin/env python3
# 导入必要的库
import argparse
import os
import sys
from time import sleep
import grpc
import traceback

# 添加P4Runtime工具库路径
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 './utils/')
)
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# 端口常量定义
SWITCH_TO_HOST_PORT = 1  # 交换机连接主机的固定端口
# 交换机之间的直接连接端口
SWITCH_TO_SWITCH_PORT_S1_S2 = 2  # s1与s2互联端口
SWITCH_TO_SWITCH_PORT_S1_S3 = 3  # s1与s3互联端口（新增）
SWITCH_TO_SWITCH_PORT_S2_S1 = 2  # s2与s1互联端口
SWITCH_TO_SWITCH_PORT_S2_S3 = 3  # s2与s3互联端口（新增）
SWITCH_TO_SWITCH_PORT_S3_S1 = 2  # s3与s1互联端口（新增）
SWITCH_TO_SWITCH_PORT_S3_S2 = 3  # s3与s2互联端口（新增）

# 隧道ID定义（每对主机双向独立，共6条隧道）
TUNNEL_ID_H1_TO_H2 = 100  # h1到h2的隧道ID
TUNNEL_ID_H2_TO_H1 = 200  # h2到h1的隧道ID
TUNNEL_ID_H1_TO_H3 = 300  # h1到h3的隧道ID（新增）
TUNNEL_ID_H3_TO_H1 = 400  # h3到h1的隧道ID（新增）
TUNNEL_ID_H2_TO_H3 = 500  # h2到h3的隧道ID（新增）
TUNNEL_ID_H3_TO_H2 = 600  # h3到h2的隧道ID（新增）

# 主机MAC地址
H1_MAC = "08:00:00:00:01:11"  # h1的MAC地址
H2_MAC = "08:00:00:00:02:22"  # h2的MAC地址
H3_MAC = "08:00:00:00:03:33"  # h3的MAC地址（新增）


def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr, transit_port):
    """
    安装隧道规则（不使用中间节点，直接在源和目标交换机之间建立隧道）

    参数：
    - p4info_helper: P4Info辅助对象
    - ingress_sw: 入口交换机（流量进入隧道的交换机）
    - egress_sw: 出口交换机（流量离开隧道的交换机）
    - tunnel_id: 隧道唯一标识ID
    - dst_eth_addr: 目标主机的MAC地址
    - dst_ip_addr: 目标主机的IP地址
    - transit_port: 入口交换机用于转发到出口交换机的端口
    """
    # 1. 入口规则：匹配目标IP，封装隧道头部
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",  # P4程序中的IPv4最长前缀匹配表
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)  # 匹配32位精确的目标IP地址
        },
        action_name="MyIngress.myTunnel_ingress",  # 执行隧道封装动作
        action_params={
            "dst_id": tunnel_id,  # 设置隧道ID
        })
    ingress_sw.WriteTableEntry(table_entry)  # 将表项写入入口交换机
    print(f"Installed ingress tunnel rule on {ingress_sw.name} (tunnel {tunnel_id})")

    # 2. 转发规则：入口交换机直接转发到出口交换机（不经过中间节点）
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",  # P4程序中的隧道精确匹配表
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id  # 匹配隧道ID
        },
        action_name="MyIngress.myTunnel_forward",  # 执行隧道转发动作
        action_params={
            "port": transit_port  # 直接转发到目标交换机的端口
        })
    ingress_sw.WriteTableEntry(table_entry)  # 将表项写入入口交换机
    print(f"Installed forward tunnel rule on {ingress_sw.name} (tunnel {tunnel_id})")

    # 3. 出口规则：解封装隧道头部，转发到目标主机
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",  # P4程序中的隧道精确匹配表
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id  # 匹配隧道ID
        },
        action_name="MyIngress.myTunnel_egress",  # 执行隧道解封装动作
        action_params={
            "dstAddr": dst_eth_addr,  # 设置目标主机MAC地址
            "port": SWITCH_TO_HOST_PORT  # 转发到主机端口（端口1）
        })
    egress_sw.WriteTableEntry(table_entry)  # 将表项写入出口交换机
    print(f"Installed egress tunnel rule on {egress_sw.name} (tunnel {tunnel_id})")


def readTableRules(p4info_helper, sw):
    """读取交换机上所有表的已安装规则并打印"""
    print(f'\n----- Reading tables rules for {sw.name} -----')
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print(f'{table_name}: ', end=' ')
            for m in entry.match:
                field_name = p4info_helper.get_match_field_name(table_name, m.field_id)
                field_value = p4info_helper.get_match_field_value(m)
                print(f'{field_name} {field_value!r} ', end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print(f'-> {action_name} ', end=' ')
            for p in action.params:
                param_name = p4info_helper.get_action_param_name(action_name, p.param_id)
                print(f'{param_name} {p.value!r} ', end=' ')
            print()


def printCounter(p4info_helper, sw, counter_name, index):
    """读取交换机上指定计数器的值，打印并返回对应格式的字符串"""
    result = ""
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            if counter_name == "MyIngress.ingressTunnelCounter":
                result = f"{sw.name}发送{counter.data.packet_count}个包，计数器编号为：{index}"
            elif counter_name == "MyIngress.egressTunnelCounter":
                result = f"{sw.name}收到{counter.data.packet_count}个包，计数器编号为：{index}"
            print(result)
    return result

def printGrpcError(e):
    """打印gRPC通信错误的详细信息"""
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print(f"({status_code.name})", end=' ')
    traceback_obj = sys.exc_info()[2]
    print(f"[{traceback_obj.tb_frame.f_code.co_filename}:{traceback_obj.tb_lineno}]")


def main(p4info_file_path, bmv2_file_path):
    """主函数：初始化连接、安装P4程序、配置隧道规则并监控计数器"""
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)
    if not os.path.exists("logs"):
        os.makedirs("logs")

    # 打开日志文件（在logs目录下）
    s1s2_log = open("logs/s1s2.txt", "a")
    s1s3_log = open("logs/s1s3.txt", "a")
    s2s3_log = open("logs/s2s3.txt", "a")

    try:
        # 创建与交换机s1的连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        # 创建与交换机s2的连接
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')

        # 创建与交换机s3的连接（新增）
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # 发送主控制器仲裁消息
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 在交换机上安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s1")

        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s2")

        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s3")

        # ================= 配置6条直接隧道规则 =================
        # 1. h1 → h2：s1直接到s2（端口2）
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2,
                         tunnel_id=TUNNEL_ID_H1_TO_H2,
                         dst_eth_addr=H2_MAC, dst_ip_addr="10.0.2.2",
                         transit_port=SWITCH_TO_SWITCH_PORT_S1_S2)

        # 2. h2 → h1：s2直接到s1（端口2）
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1,
                         tunnel_id=TUNNEL_ID_H2_TO_H1,
                         dst_eth_addr=H1_MAC, dst_ip_addr="10.0.1.1",
                         transit_port=SWITCH_TO_SWITCH_PORT_S2_S1)

        # 3. h1 → h3：s1直接到s3（端口3）（新增）
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3,
                         tunnel_id=TUNNEL_ID_H1_TO_H3,
                         dst_eth_addr=H3_MAC, dst_ip_addr="10.0.3.3",
                         transit_port=SWITCH_TO_SWITCH_PORT_S1_S3)

        # 4. h3 → h1：s3直接到s1（端口2）（新增）
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1,
                         tunnel_id=TUNNEL_ID_H3_TO_H1,
                         dst_eth_addr=H1_MAC, dst_ip_addr="10.0.1.1",
                         transit_port=SWITCH_TO_SWITCH_PORT_S3_S1)

        # 5. h2 → h3：s2直接到s3（端口3）（新增）
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3,
                         tunnel_id=TUNNEL_ID_H2_TO_H3,
                         dst_eth_addr=H3_MAC, dst_ip_addr="10.0.3.3",
                         transit_port=SWITCH_TO_SWITCH_PORT_S2_S3)

        # 6. h3 → h2：s3直接到s2（端口3）（新增）
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2,
                         tunnel_id=TUNNEL_ID_H3_TO_H2,
                         dst_eth_addr=H2_MAC, dst_ip_addr="10.0.2.2",
                         transit_port=SWITCH_TO_SWITCH_PORT_S3_S2)

        # 读取所有交换机表项
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        # 循环监控隧道计数器
        while True:
            sleep(2)
            print('\n----- Reading tunnel counters -----')

            print("----- s1 -> s2 -----")
            s1s2_log.write("----- s1 -> s2 -----\n")
            s1_send = printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", TUNNEL_ID_H1_TO_H2)
            s1s2_log.write(s1_send + "\n")
            s2_receive = printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", TUNNEL_ID_H1_TO_H2)
            s1s2_log.write(s2_receive + "\n")

            print("----- s2 -> s1 -----")
            s1s2_log.write("----- s2 -> s1 -----\n")
            s2_send = printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", TUNNEL_ID_H2_TO_H1)
            s1s2_log.write(s2_send + "\n")
            s1_receive = printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", TUNNEL_ID_H2_TO_H1)
            s1s2_log.write(s1_receive + "\n")

            # h1↔h3计数器（s1与s3之间）
            print("----- s1 -> s3 -----")
            s1s3_log.write("----- s1 -> s3 -----\n")
            s1_to_s3_send = printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", TUNNEL_ID_H1_TO_H3)
            s1s3_log.write(s1_to_s3_send + "\n")
            s3_from_s1_receive = printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", TUNNEL_ID_H1_TO_H3)
            s1s3_log.write(s3_from_s1_receive + "\n")

            print("----- s3 -> s1 -----")
            s1s3_log.write("----- s3 -> s1 -----\n")
            s3_to_s1_send = printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", TUNNEL_ID_H3_TO_H1)
            s1s3_log.write(s3_to_s1_send + "\n")
            s1_from_s3_receive = printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", TUNNEL_ID_H3_TO_H1)
            s1s3_log.write(s1_from_s3_receive + "\n")

            # h2↔h3计数器（s2与s3之间）
            print("----- s2 -> s3 -----")
            s2s3_log.write("----- s2 -> s3 -----\n")
            s2_to_s3_send = printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", TUNNEL_ID_H2_TO_H3)
            s2s3_log.write(s2_to_s3_send + "\n")
            s3_from_s2_receive = printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", TUNNEL_ID_H2_TO_H3)
            s2s3_log.write(s3_from_s2_receive + "\n")

            print("----- s3 -> s2 -----")
            s2s3_log.write("----- s3 -> s2 -----\n")
            s3_to_s2_send = printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", TUNNEL_ID_H3_TO_H2)
            s2s3_log.write(s3_to_s2_send + "\n")
            s2_from_s3_receive = printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", TUNNEL_ID_H3_TO_H2)
            s2s3_log.write(s2_from_s3_receive + "\n")


            s1s2_log.flush()
            s1s3_log.flush()
            s2s3_log.flush()

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"Unexpected error: {str(e)}")
        traceback.print_exc()
    finally:
        # 关闭日志文件
        s1s2_log.close()
        s1s3_log.close()
        s2s3_log.close()

    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
